package com.lenovo.spider.support.cookie;

import cn.hutool.core.collection.CollUtil;
import lombok.val;
import okhttp3.Cookie;
import okhttp3.CookieJar;
import okhttp3.HttpUrl;

import java.util.List;
import java.util.Map;
import java.util.Objects;

import static java.util.stream.Collectors.toList;

/**
 * @author <a href="mailto:liuhan3@lenovo.com">HanL(liuhan3)</a>
 * on 2018/3/29 17:22
 */
public class SpiderCookieJar implements CookieJar {

    private final CookieStore cookieStore;

    public SpiderCookieJar(CookieStore cookieStore) {
        this.cookieStore = cookieStore;
    }

    /**
     * 存储cookie
     *
     * @param url
     * @param cookies
     */
    @Override
    public void saveFromResponse(HttpUrl url, List<Cookie> cookies) {
        cookieStore.appendCookie(url, cookies);
    }

    /**
     * 加载cookie
     *
     * @param url
     * @return
     */
    @Override
    public List<Cookie> loadForRequest(HttpUrl url) {
        return cookieStore.getCookies(url);
    }

    /**
     * 添加自定义cookie
     *
     * @param url
     * @param cookies
     */
    public void addCustomCookies(String url, Map<String, String> cookies) {
        if (CollUtil.isEmpty(cookies)) return;
        val httpUrl = HttpUrl.parse(url);
        val domain = Objects.requireNonNull(httpUrl).host();
        List<Cookie> cookieList = cookies.entrySet().stream()
                .map(entity -> new Cookie.Builder()
                        .domain(domain)
                        .path("/")
                        .name(entity.getKey())
                        .value(entity.getValue())
                        .build())
                .collect(toList());
        saveFromResponse(httpUrl, cookieList);
    }

}
